{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RWTH Plots with Matplotlib\n",
    "`rwth_nb.plots.mpl_decorations` extends Matplotlib to some useful functionality explained below\n",
    "\n",
    "1. [Simple Plots](#Simple-Plots)\n",
    "    1. [Graph Plot](#Graph-Plot)\n",
    "    2. [Stem Plot](#Stem-Plot)\n",
    "    3. [Multiple Plots](#Multiple-Plots)\n",
    "    4. [Updating Plots](#Updating-Plots)\n",
    "2. [Annotations](#Annotations)\n",
    "    1. [Ticks](#Ticks)\n",
    "    2. [Distances](#Distances)\n",
    "3. [Misc](#Misc)\n",
    "    1. [Signal processing](#Signal-processing)\n",
    "        1. [Dirac Impulses](#Dirac-impulses)\n",
    "        2. [Pole-Zero Diagrams](#Pole-Zero-Diagrams)\n",
    "        \n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple Plots\n",
    "\n",
    "### Graph Plot\n",
    "\n",
    "See also:\n",
    "* [Matplotlib pyplot tutorial](https://matplotlib.org/3.2.1/tutorials/introductory/pyplot.html)\n",
    "* [Axis Styles](#Axis-Styles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rwth_nb.plots.mpl_decorations as rwth_plt\n",
    "\n",
    "# numpy array\n",
    "x = np.linspace(-5,5)\n",
    "\n",
    "# create figure\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "# call axis and grid for beautification\n",
    "rwth_plt.axis(ax); \n",
    "rwth_plt.grid(ax)\n",
    "\n",
    "# plot to axis while color is in RWTH Colors\n",
    "ax.plot(x, x**2, 'rwth:blue');  \n",
    "\n",
    "# set x- and y-labels\n",
    "ax.set_xlabel(r'$\\rightarrow x$'); \n",
    "ax.set_ylabel(r'$\\uparrow f(x)=x^2$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where\n",
    "* `rwth_plt.axis(ax)` relocates the axes spines,\n",
    "* `rwth_plt.grid(ax)` displays an additional grid."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Axis Styles\n",
    "Relocate axes spines.\n",
    "```\n",
    "rwth_plt.axis(ax)\n",
    "```\n",
    "\n",
    "#### Grid\n",
    "Displays grid.\n",
    "```\n",
    "rwth_plt.grid(ax)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stem Plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See also:\n",
    "\n",
    "* [Matplotlib stem examples](https://matplotlib.org/3.2.1/gallery/lines_bars_and_markers/stem_plot.html)\n",
    "* [Axis Styles](#Axis-Styles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rwth_nb.plots.mpl_decorations as rwth_plt\n",
    "\n",
    "n = np.linspace(-5, 5, 11)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax)\n",
    "ax.set_xlabel(r'$\\rightarrow n$'); \n",
    "ax.set_ylabel(r'$\\uparrow f(n)=n$');\n",
    "\n",
    "# stem plot\n",
    "rwth_plt.stem(ax, n, n, 'rwth:blue');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiple Plots\n",
    "\n",
    "Plots can be combined by using `matplotlib.pyplot.subplots()`, where arguments are unpacked from `rwth_plt.landscape` for 16/9 landscape view.\n",
    "Optionally, own figure sizes can be defined.\n",
    "\n",
    "See also: [Matplotlib subplots](https://matplotlib.org/3.2.1/api/_as_gen/matplotlib.pyplot.subplots.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rwth_nb.plots.mpl_decorations as rwth_plt\n",
    "\n",
    "t = np.linspace(-5, 5, 10001)\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, **rwth_plt.landscape)\n",
    "\n",
    "ax = axs[0]\n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax)\n",
    "ax.plot(t, t**2, 'rwth:blue')\n",
    "\n",
    "ax = axs[1]\n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax)\n",
    "ax.plot(t, t**3, 'rwth:red');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Updating Plots\n",
    "Plots can be dynamically updated for more efficient usage.  \n",
    "This is essential when using [Widgets](RWTH%20Widgets.ipynb).  \n",
    "\n",
    "#### Graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rwth_nb.plots.mpl_decorations as rwth_plt\n",
    "\n",
    "from ipywidgets import widgets\n",
    "\n",
    "# initial plot\n",
    "# create figure\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "# call axis and grid for beautification\n",
    "rwth_plt.axis(ax);\n",
    "rwth_plt.grid(ax)\n",
    "\n",
    "# plot to axis\n",
    "# this time: save plotted lines to a variable\n",
    "x = np.linspace(-5, 5, 10001)\n",
    "n = 1\n",
    "\n",
    "line, = ax.plot(x, x**n, 'rwth:blue');\n",
    "\n",
    "# set x- and y-labels\n",
    "ax.set_xlabel(r'$\\rightarrow x$'); \n",
    "ax.set_ylabel(r'$\\uparrow f(x)=x^{}$'.format(n));\n",
    "\n",
    "# set up widget for updating n (see RWTH Widgets.ipynb) \n",
    "@widgets.interact(n=widgets.IntSlider(min=1, max=6, step=1, description='$n$', style=rwth_plt.wdgtl_style))\n",
    "def update_n(n):\n",
    "    # updating plot\n",
    "    # change lines' y-data\n",
    "    line.set_ydata(x**n)\n",
    "    \n",
    "    # change label\n",
    "    ax.set_ylabel(r'$\\uparrow f(x)=x^{}$'.format(n));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Stems\n",
    "\n",
    "Stem plots are stored in so called containers. These can be manipulated efficiently using\n",
    "\n",
    "```\n",
    "stem_set_data(container, x, y)\n",
    "stem_set_xdata(container, x)\n",
    "stem_set_ydata(container, y)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rwth_nb.plots.mpl_decorations as rwth_plt\n",
    "\n",
    "from ipywidgets import widgets\n",
    "\n",
    "a = .5\n",
    "n = np.linspace(-5, 5, 11)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax)\n",
    "\n",
    "ax.set_xlabel(r'$\\rightarrow n$'); \n",
    "ax.set_ylabel(r'$\\uparrow f(n)={}\\cdot n$'.format(a));\n",
    "\n",
    "# initial stem plot\n",
    "stem_container = rwth_plt.stem(ax, n, n, 'rwth:blue');\n",
    "\n",
    "# set up widget for updating n (see RWTH Widgets.ipynb) \n",
    "@widgets.interact(a=widgets.FloatSlider(min=-1, max=1, step=.1, value=.5, description='$a$', style=rwth_plt.wdgtl_style))\n",
    "def update_a(a):\n",
    "    \n",
    "    # change stem data\n",
    "    rwth_plt.stem_set_ydata(stem_container, a*n)\n",
    "    \n",
    "    # change label\n",
    "    ax.set_ylabel(r'$\\uparrow f(x)={}\\cdot n$'.format(a));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ticks\n",
    "`annotate_xtick`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "x0 = 2\n",
    "\n",
    "fig,ax = plt.subplots(); ax.plot(x, (x-x0)**2, 'rwth:blue');\n",
    "ax.set_xlabel(r'$\\rightarrow x$'); ax.set_ylabel(r'$\\uparrow f(x)=(x-x_0)^2$')\n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax);\n",
    "\n",
    "rwth_plt.annotate_xtick(ax, r'$x_0$', x0, x0**2 );"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`annotate_ytick`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = 2\n",
    "\n",
    "fig,ax = plt.subplots(); ax.plot(x, (x-x0)**2, 'rwth:blue');\n",
    "ax.set_xlabel(r'$\\rightarrow x$'); ax.set_ylabel(r'$\\uparrow f(x)=(x-x_0)^2$')\n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax);\n",
    "\n",
    "rwth_plt.annotate_ytick(ax, r'$y_0$', -1, 4 );"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distances\n",
    "`annotate_distance`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Delta = 3\n",
    "\n",
    "fig,ax = plt.subplots(); ax.plot(x, np.exp(-(x/np.sqrt(Delta))**2), 'rwth:blue');\n",
    "ax.set_xlabel(r'$\\rightarrow x$'); \n",
    "rwth_plt.axis(ax); rwth_plt.grid(ax);\n",
    "\n",
    "rwth_plt.annotate_distance(ax, r'$\\Delta$', [-Delta/2,.5], [Delta/2,.5]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Misc\n",
    "\n",
    "### Signal processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dirac impulses\n",
    "`plot_dirac`, `dirac_weights`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up plot\n",
    "fig,ax = plt.subplots(1,1);\n",
    "ax.set_xlim([-2.75, 2.75]); ax.set_ylim([0, 2.4]); rwth_plt.grid(ax); rwth_plt.axis(ax);\n",
    "\n",
    "# dirac x positions\n",
    "dirac_x = [-1, 0, 1]\n",
    "\n",
    "# dirac weights\n",
    "dirac_weights = [1, 2, 1]\n",
    "\n",
    "# plot diracs\n",
    "rwth_plt.plot_dirac(ax, dirac_x, dirac_weights, 'rwth:blue');\n",
    "\n",
    "# show weights\n",
    "rwth_plt.dirac_weights(ax, dirac_x, dirac_weights, dirac_weights, color='rwth:black')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pole-Zero Diagrams\n",
    "\n",
    "`plot_lroc` Laplace Domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# poles (pp) and zeros (pz)\n",
    "pp = np.array([-1, 1+1j]); pz = np.array([0, 3-2j])\n",
    "\n",
    "# poles' and zeros' orders\n",
    "ord_p = np.array([1, 2]); ord_z = np.array([1, 1])\n",
    "\n",
    "# region of convergence\n",
    "roc = np.array([np.max(np.real(pp)), np.inf])\n",
    "\n",
    "# set up plot\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlabel(r'$\\rightarrow \\mathrm{Re}$'); ax.set_ylabel(r'$\\uparrow \\mathrm{Im}$');\n",
    "ax.set_xlim(-2.5, 3.5); ax.set_ylim(-5, 5); rwth_plt.grid(ax); rwth_plt.axis(ax); ax.set_title('Pole-Zero Diagram');\n",
    "\n",
    "# plot poles and zeros\n",
    "ax.plot(np.real(pp), np.imag(pp), **rwth_plt.style_poles); ax.plot(np.real(pp), -np.imag(pp), **rwth_plt.style_poles);\n",
    "ax.plot(np.real(pz), np.imag(pz), **rwth_plt.style_zeros); ax.plot(np.real(pz), -np.imag(pz), **rwth_plt.style_zeros);\n",
    "\n",
    "# show poles' and zeros' orders\n",
    "rwth_plt.annotate_order(ax, pp, ord_p)\n",
    "rwth_plt.annotate_order(ax, pz, ord_z)\n",
    "\n",
    "# show S_0\n",
    "S_0 = 1\n",
    "ax.text(2, 4, r'$S_0 =$ ' + str(S_0), fontsize=12, bbox=rwth_plt.wbbox)\n",
    "\n",
    "# plot region of convergence\n",
    "rwth_plt.plot_lroc(ax, roc);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`plot_zroc` $z$-Domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# poles (zp) and zeros (zz)\n",
    "zp = np.array([0.5, 1+1j]); zz = np.array([0])\n",
    "\n",
    "# poles' and zeros' orders\n",
    "ord_p = np.array([1, 2]); ord_z = np.array([1])\n",
    "\n",
    "# region of convergence\n",
    "roc = np.array([0.5, 1+1j])\n",
    "\n",
    "# set up plot\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlabel(r'$\\rightarrow \\mathrm{Re}$'); ax.set_ylabel(r'$\\uparrow \\mathrm{Im}$');\n",
    "rwth_plt.grid(ax); rwth_plt.axis(ax); ax.set_title('Pole-Zero Diagram');\n",
    "\n",
    "# square plot (!)\n",
    "ax.set_aspect('equal')\n",
    "\n",
    "# plot poles and zeros\n",
    "ax.plot(np.real(zp), np.imag(zp), **rwth_plt.style_poles); ax.plot(np.real(zp), -np.imag(zp), **rwth_plt.style_poles);\n",
    "ax.plot(np.real(zz), np.imag(zz), **rwth_plt.style_zeros); ax.plot(np.real(zz), -np.imag(zz), **rwth_plt.style_zeros);\n",
    "\n",
    "# show orders\n",
    "rwth_plt.annotate_order(ax, zp, ord_p)\n",
    "rwth_plt.annotate_order(ax, zz, ord_z)\n",
    "\n",
    "# plot region of convergence\n",
    "rwth_plt.plot_zroc(ax, roc);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "This code is licensed under the [MIT license](https://opensource.org/licenses/MIT)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}